// This file implements compilation of nodes that contain binding
// pairs (i.e. `let` and attribute sets).
//
// Bindings are fairly complex: There can be nested attribute sets
// (e.g. `a.b = 42`), bindings can recursively refer to each other,
// there are different scopes in which `inherit` operates and so on.
//
// Due to this the compilation of bindings is likely the most complex
// part of Nix evaluation (apart from thunk states).
package compiler

import (
	"errors"
	"fmt"

	"tvl.fyi/tvix/eval/internal/index"
	"tvl.fyi/tvix/eval/internal/op"
	p "tvl.fyi/tvix/eval/parser"
	"tvl.fyi/tvix/eval/value"
)

// staticIdent returns a statically known identifier from a node if possible.
func (c *compiler) staticIdent(node *p.Node) string {
	switch node.Type {
	case p.IDNode:
		return c.parser.TokenString(node.Tokens[0])

	case p.StringNode:
		if len(node.Nodes) == 0 {
			return ""
		}

		textNode := node.Nodes[0]
		// TODO: decode string
		return c.parser.TokenString(textNode.Tokens[0])

	case p.InterpNode:
		// Only dynamic keys that actually only contain a static string
		// can be used here.
		if node.Nodes[0].Type == p.StringNode {
			return c.staticIdent(node.Nodes[0])
		}
	}

	return ""
}

// firstPathIdent returns the first identifier of an attribute path, and the
// remaining path.
func (c *compiler) firstPathIdent(path []*p.Node) (string, []*p.Node) {
	if len(path) == 0 {
		return "", nil
	}

	return c.staticIdent(path[0]), path[1:]
}

// compileAttrs emits bytecode for non-recursive attribute sets.
//
// Nodes can either be bind (e.g. `a = b`) or inherit.
func (c *compiler) compileAttrs(slot index.Local, nodes []*p.Node) {
	if len(nodes) == 0 {
		c.pushConstant(slot, value.NixAttrs{})
		return
	}

	// Open a scope to track the positions of the temporaries used by the
	// `op.Attrs` instruction.
	c.unit().beginScope()

	var bindings, inheritFroms []*p.Node

	// First pass: Plain inherits are compiled right away, other binding
	// types are collected for the next passes.
	for _, node := range nodes {
		switch node.Type {
		case p.InheritNode:
			c.error(fmt.Errorf("TODO: plain inherits not yet supported"))
		case p.BindNode:
			bindings = append(bindings, node)
		case p.InheritFromNode:
			inheritFroms = append(inheritFroms, node)
		default:
			c.error(fmt.Errorf("tvix bug: unknown attrs node type %q", node.Type))
		}
	}

	// Second pass: Binding stack slots are declared, and nested attribute
	// sets are collected into new data structures. For non-recursive sets,
	// all stack slots are phantoms.
	for _, i := range inheritFroms {
		c.assertNodeType(i, p.InheritFromNode)
		c.error(fmt.Errorf("TODO: namespaced inherits are not yet supported"))
	}

	tracked := make([]trackedBinding, 0, len(bindings))
	staticKeys := make(map[string]int, len(bindings))

	for _, b := range bindings {
		c.assertNodeType(b, p.BindNode)

		// bind.attrpath.<children> = attrpath nodes
		key, rem := c.firstPathIdent(b.Nodes[0].Nodes)

		if key == "" { // dynamic keys are never merged
			tracked = append(tracked, newBinding(
				trackedKey{c.unit().declarePhantom(), "", b.Nodes[0].Nodes[0]},
				c.unit().declarePhantom(),
				rem, b.Nodes[1],
			))

			continue
		}

		if idx, found := staticKeys[key]; found {
			c.error(tracked[idx].merge(rem, b.Nodes[1]))
			continue
		}

		idx := len(tracked)
		tracked = append(tracked, newBinding(
			trackedKey{c.unit().declarePhantom(), key, nil},
			c.unit().declarePhantom(),
			rem, b.Nodes[1],
		))

		staticKeys[key] = idx
	}

	// Third pass: Actually compile all values into the correct locations.
	for _, t := range tracked {
		// key:
		if t.key.static != "" {
			c.pushConstant(t.key.slot, value.String(t.key.static))
		} else {
			c.compileNode(t.key.slot, t.key.dynamic)
		}

		// value:
		c.compileNode(t.valueSlot, t.value)
	}

	// TODO: fix pair count for inherits
	c.unit().chunk.PushOpWithUint(op.Attrs, uint64(len(tracked)))

	// Remove the temporary scope, but do not emit any additional cleanup
	// (op.Attrs consumes all of these locals).
	c.unit().endScope()
}

// compileLet emits bytecode for a `let ... in ...` expression.
func (c *compiler) compileLet(slot index.Local, node *p.Node) {
	c.unit().beginScope()

	if len(node.Nodes[0].Nodes) == 0 {
		c.compileNode(slot, node.Nodes[1])
		return
	}

	var bindings []*p.Node

	// First pass (split)
	for _, node := range node.Nodes[0].Nodes { // attrs -> binds -> [bind]
		switch node.Type {
		case p.BindNode:
			bindings = append(bindings, node)
		default:
			c.error(fmt.Errorf("unsupported node in let: %q", node.Type))
		}
	}

	// Second pass (declare)

	tracked := make([]trackedBinding, 0, len(bindings))
	staticKeys := make(map[string]int, len(bindings))

	for _, b := range bindings {
		c.assertNodeType(b, p.BindNode)

		// bind.attrpath.<children> = attrpath nodes
		key, rem := c.firstPathIdent(b.Nodes[0].Nodes)
		if key == "" {
			c.error(fmt.Errorf("dynamic keys not allowed in let-binding"))
			continue
		}

		if idx, found := staticKeys[key]; found {
			c.error(tracked[idx].merge(rem, b.Nodes[1]))
			continue
		}

		idx := len(tracked)
		slot, _ := c.unit().declareLocal(key) // TODO: shadowed?

		tracked = append(tracked, newBinding(
			trackedKey{keyNotEmitted, key, nil},
			slot, rem, b.Nodes[1],
		))
		staticKeys[key] = idx
	}

	// Third pass (compile)
	for _, t := range tracked {
		c.thunk(t.valueSlot, t.value)
	}

	// Fourth pass: Emit capture instructions if necessary.
	for _, t := range tracked {
		l := c.unit().locals[t.valueSlot]
		if len(l.captures) == 0 {
			continue
		}

		args := make([]uint64, len(l.captures)+1)

		// Store the target index to write captures to as the first arg.
		args[0] = uint64(c.unit().stackIndex(t.valueSlot))

		for i, capture := range l.captures {
			// For captured locals, calculate the stack slot as the argument.
			if capturedLocal, ok := capture.CompileTimeLocalIndex(); ok {
				args[i+1] = uint64(c.unit().stackIndex(capturedLocal))
				continue
			}

			// For captured captures, emit the parent's capture index.
			capturedCapture, _ := capture.CaptureIndex()
			args[i+1] = uint64(capturedCapture)
		}

		c.unit().chunk.PushOpWithUintArgs(op.Capture, args...)
	}

	// Body ...
	c.compileNode(slot, node.Nodes[1])

	pops := c.unit().endScope()
	c.unit().chunk.PushOpWithUint(op.CloseScope, uint64(pops))
}

func (c *compiler) compileID(slot index.Local, node *p.Node) {
	// only non-recursive, simple identifier access is supported so far
	ident := c.parser.TokenString(node.Tokens[0])
	pos := c.unit().resolveLocal(ident)

	if pos != index.Unknown {
		if idx, local := pos.CompileTimeLocalIndex(); local {
			stackIdx := c.unit().stackIndex(idx)
			c.unit().chunk.PushOpWithUint(op.GetLocal, uint64(stackIdx))
			return
		}

		if idx, capture := pos.CaptureIndex(); capture {
			c.unit().chunk.PushOpWithUint(op.GetCapture, uint64(idx))
			return
		}
	}

	// try globals
	if v, ok := globals()[ident]; ok {
		c.pushConstant(slot, v)
		return
	}

	c.error(fmt.Errorf("undefined variable '%s'", ident))
}

// Types for tracking bindings before compiling their values, and
// merging nested attribute sets.

const keyNotEmitted index.Local = -2

type trackedKey struct {
	// slot holds the local key slot in which the key is being
	// compiled. The special value keyNotEmitted signifies that
	// this key belongs to a `let`-binding and does not have a
	// stack slot.
	slot index.Local

	// static holds the statically known key of the attribute, if
	// present.
	static string

	// dynamic holds a pointer to the AST node that yields the key
	// at runtime. If this is set the key can not be statically
	// known.
	dynamic *p.Node
}

func (t *trackedKey) String() string {
	if t.static != "" {
		return fmt.Sprintf("key %q", t.static)
	}

	return "dynamic key"
}

type trackedBinding struct {
	key       trackedKey
	valueSlot index.Local
	value     *p.Node
}

func bindNode(path []*p.Node, value *p.Node) *p.Node {
	return &p.Node{
		Type: p.BindNode,
		Nodes: []*p.Node{
			// first element is the nested path, second is the value
			{
				Type:  p.AttrPathNode,
				Nodes: path,
			},
			value,
		},
	}
}

func newBinding(key trackedKey, slot index.Local, path []*p.Node, value *p.Node) trackedBinding {
	b := trackedBinding{
		key:       key,
		valueSlot: slot,
	}

	if len(path) == 0 {
		b.value = value
		return b
	}

	b.value = &p.Node{
		Type:  p.SetNode,
		Nodes: []*p.Node{bindNode(path, value)},
	}

	return b
}

func (t *trackedBinding) merge(path []*p.Node, value *p.Node) error {
	if len(path) == 0 {
		return errors.New(t.key.String() + " already defined")
	}

	if !(t.value.Type == p.SetNode || t.value.Type == p.RecSetNode) {
		return fmt.Errorf("can only merge attribute sets, but node is a %s", t.value.Type)
	}

	t.value.Nodes = append(t.value.Nodes, bindNode(path, value))
	return nil
}
